import cv2
import numpy as np
import random
import os
from datetime import datetime, timedelta
from PIL import Image, ImageDraw, ImageFont
import shutil


def generate_random_date():
    # 定义日期范围，这里假设从 2000 年 1 月 1 日到 2030 年 12 月 31 日
    start_date = datetime(2000, 1, 1)
    end_date = datetime(2030, 12, 31)
    # 计算日期范围的总秒数
    total_seconds = (end_date - start_date).total_seconds()
    # 随机选择一个秒数
    random_seconds = random.randint(0, int(total_seconds))
    # 生成随机日期
    random_date = start_date + timedelta(seconds=random_seconds)
    # 格式化日期为指定格式
    return random_date.strftime("%Y年%m月%d日 %H:%M:%S")


def add_text_in_polygons(image_path, yolo_txt_path):
    # 读取图片
    image = cv2.imread(image_path)
    if image is None:
        print(f"无法读取图片: {image_path}")
        return

    height, width = image.shape[:2]

    try:
        with open(yolo_txt_path, 'r') as f:
            lines = f.readlines()
    except FileNotFoundError:
        print(f"未找到标注文件: {yolo_txt_path}")
        return

    # 将 OpenCV 图像转换为 PIL 图像
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    pil_image = Image.fromarray(image)
    draw = ImageDraw.Draw(pil_image)
    # 指定中文字体文件路径，需要根据实际情况修改
    font = ImageFont.truetype("simhei.ttf", 20)

    for line in lines:
        parts = line.strip().split()
        if len(parts) == 0:
            continue
        class_id = int(parts[0])
        polygon_points = [float(x) for x in parts[1:]]

        # 将归一化坐标转换为实际坐标
        points = []
        for i in range(0, len(polygon_points), 2):
            x = int(polygon_points[i] * width)
            y = int(polygon_points[i + 1] * height)
            points.append([x, y])
        points = np.array(points, dtype=np.int32)

        # 计算多边形的边界框
        x, y, w, h = cv2.boundingRect(points)
        # 在多边形内部随机选择一个位置
        while True:
            random_x = random.randint(x, x + w)
            random_y = random.randint(y, y + h)
            # 检查随机点是否在多边形内部
            if cv2.pointPolygonTest(points, (random_x, random_y), False) >= 0:
                break

        # 生成随机日期
        random_date = generate_random_date()
        # 随机选择文字颜色为纯白或纯黑
        text_color = (255, 255, 255) if random.choice([True, False]) else (0, 0, 0)
        # 在随机位置添加文字
        draw.text((random_x, random_y), random_date, font=font, fill=text_color)

    # 将 PIL 图像转换回 OpenCV 图像
    image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
    return image


def run(image_dir, yolo_txt_dir, output_dir, output_labels_dir):
    # 遍历图片和标注文件
    for filename in os.listdir(image_dir):
        if filename.endswith(('.png', '.jpg', '.jpeg')):
            image_path = os.path.join(image_dir, filename)
            base_name = os.path.splitext(filename)[0]
            yolo_txt_path = os.path.join(yolo_txt_dir, f"{base_name}.txt")
            output_image = add_text_in_polygons(image_path, yolo_txt_path)
            if output_image is not None:
                file_extension = os.path.splitext(filename)[1]
                output_filename = f"{base_name}{file_extension}"
                output_path = os.path.join(output_dir, output_filename)
                cv2.imwrite(output_path, output_image)
                shutil.copy2(yolo_txt_path, os.path.join(output_labels_dir, f"{base_name}.txt"))


if __name__ == '__main__':
    image_dir = "../yolo/train/images"
    yolo_txt_dir = "../yolo/train/labels"
    output_dir = "../test/train/images"
    output_labels_dir = "../test/train/labels"
    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(output_labels_dir, exist_ok=True)
    run(image_dir, yolo_txt_dir, output_dir, output_labels_dir)

    image_dir = "../yolo/val/images"
    yolo_txt_dir = "../yolo/val/labels"
    output_dir = "../test/val/images"
    output_labels_dir = "../test/val/labels"
    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(output_labels_dir, exist_ok=True)
    run(image_dir, yolo_txt_dir, output_dir, output_labels_dir)
